package com.example.gpu;

import jcuda.*;
import jcuda.driver.CUdeviceptr;
import jcuda.jcublas.JCublas;
import jcuda.jcublas.cublasHandle;
import jcuda.runtime.cudaDeviceProp;

import static jcuda.jcublas.JCublas2.*;
import static jcuda.jcublas.cublasOperation.CUBLAS_OP_N;
import static jcuda.runtime.JCuda.*;

public class JCublasMatrixMultiplyExample {

    public static void main(String[] args) {
        // 初始化JCublas
        JCublas.cublasInit();
        
        // 创建cublas句柄
        cublasHandle handle = new cublasHandle();
        cublasCreate(handle);
        
        // 矩阵维度
        int m = 1000; // A和C的行数
        int n = 1000; // B和C的列数
        int k = 1000; // A的列数和B的行数
        
        // 标量alpha和beta用于矩阵运算
        float alpha = 1.0f;
        float beta = 0.0f;
        
        // 为矩阵分配主机内存
        float[] h_A = new float[m * k];
        float[] h_B = new float[k * n];
        float[] h_C = new float[m * n];
        
        // 初始化矩阵数据
        for (int i = 0; i < m * k; i++) {
            h_A[i] = (float)Math.random();
        }
        for (int i = 0; i < k * n; i++) {
            h_B[i] = (float)Math.random();
        }
        
        // 为矩阵分配设备内存
        cudaDeviceProp prop = new cudaDeviceProp();
        cudaGetDeviceProperties(prop, 0);
        
        CUdeviceptr d_A = new CUdeviceptr();
        CUdeviceptr d_B = new CUdeviceptr();
        CUdeviceptr d_C = new CUdeviceptr();
        
        cudaMalloc(d_A, m * k * Sizeof.FLOAT);
        cudaMalloc(d_B, k * n * Sizeof.FLOAT);
        cudaMalloc(d_C, m * n * Sizeof.FLOAT);
        
        // 将数据从主机复制到设备
        cublasSetMatrix(m, k, Sizeof.FLOAT, Pointer.to(h_A), m, d_A, m);
        cublasSetMatrix(k, n, Sizeof.FLOAT, Pointer.to(h_B), k, d_B, k);
        
        // 计时开始
        long startTime = System.currentTimeMillis();
        
        // 执行矩阵乘法: C = alpha * A * B + beta * C
        // 使用cublasSgemm函数，参数分别为:
        // 1. 句柄
        // 2. 矩阵A的转置模式
        // 3. 矩阵B的转置模式
        // 4. 矩阵维度参数
        // 5. 标量alpha
        // 6. 矩阵A的设备指针
        // 7. 矩阵A的leading dimension
        // 8. 矩阵B的设备指针
        // 9. 矩阵B的leading dimension
        // 10. 标量beta
        // 11. 矩阵C的设备指针
        // 12. 矩阵C的leading dimension
        cublasSgemm(
            handle, 
            CUBLAS_OP_N, CUBLAS_OP_N, 
            m, n, k, 
            Pointer.to(new float[]{alpha}), 
            d_A, m, 
            d_B, k, 
            Pointer.to(new float[]{beta}), 
            d_C, m
        );
        
        // 确保计算完成
        cudaDeviceSynchronize();
        
        // 计时结束
        long endTime = System.currentTimeMillis();
        
        // 将结果从设备复制回主机
        cublasGetMatrix(m, n, Sizeof.FLOAT, d_C, m, Pointer.to(h_C), m);
        
        System.out.println("GPU矩阵乘法完成");
        System.out.println("矩阵维度: " + m + "x" + k + " * " + k + "x" + n + " = " + m + "x" + n);
        System.out.println("GPU计算耗时: " + (endTime - startTime) + " ms");
        
        // 在CPU上执行相同的矩阵乘法以进行性能比较
        startTime = System.currentTimeMillis();
        float[] h_C_cpu = matrixMultiplyCPU(h_A, h_B, m, k, n);
        endTime = System.currentTimeMillis();
        
        System.out.println("CPU计算耗时: " + (endTime - startTime) + " ms");
        
        // 验证结果
        boolean passed = verifyResults(h_C, h_C_cpu, m * n);
        System.out.println("计算结果验证: " + (passed ? "通过" : "失败"));
        
        // 释放设备内存
        cudaFree(d_A);
        cudaFree(d_B);
        cudaFree(d_C);
        
        // 销毁句柄
        cublasDestroy(handle);
        
        // 关闭JCublas
        JCublas.cublasShutdown();
    }
    
    /**
     * 在CPU上执行矩阵乘法
     */
    private static float[] matrixMultiplyCPU(float[] A, float[] B, int m, int k, int n) {
        float[] C = new float[m * n];
        
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                float sum = 0.0f;
                for (int p = 0; p < k; p++) {
                    sum += A[i * k + p] * B[p * n + j];
                }
                C[i * n + j] = sum;
            }
        }
        
        return C;
    }
    
    /**
     * 验证GPU和CPU计算结果是否一致
     */
    private static boolean verifyResults(float[] gpuResult, float[] cpuResult, int size) {
        float epsilon = 1e-4f; // 允许的误差范围
        
        for (int i = 0; i < size; i++) {
            if (Math.abs(gpuResult[i] - cpuResult[i]) > epsilon) {
                System.out.println("结果不匹配在索引 " + i + ": GPU = " + gpuResult[i] + ", CPU = " + cpuResult[i]);
                return false;
            }
        }
        
        return true;
    }
}